from toolz.curried import get

from paddle.fluid import layers as L
from helm.static.engine.metrics import Average


class TopKAccuracy(Average):

    def __init__(self, name=None, k=5):
        self.k = k
        name = self._init_name(name)
        super().__init__(name)

    def append_op(self, output):
        y_pred, y_true = get(["y_pred", "y_true"], output)
        acc = L.accuracy(y_pred, y_true, k=self.k)
        return acc

    def _init_name(self, name):
        if name is None:
            name = "acc" + str(self.k)
        return name


class Accuracy(Average):

    def __init__(self, name="acc"):
        super().__init__(name)

    def append_op(self, output):
        y_pred, y_true = get(["y_pred", "y_true"], output)
        if len(y_true.shape) == 2 and y_true.shape[1] != 1:
            y_true = L.argmax(y_true, axis=1)
            y_true = L.unsqueeze(y_true, 1)
        acc = L.accuracy(y_pred, y_true)
        return acc